package main

import (
	"context"
	"fmt"
	"math/rand"
	"time"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

const (
	schedulerName = "custom-scheduler"
)

func main() {
	rand.Seed(time.Now().UnixNano())

	// 1. 创建 Kubernetes 客户端
	config, err := buildConfig()
	if err != nil {
		panic(err.Error())
	}

	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		panic(err.Error())
	}

	// 2. 监听 Pod 事件
	watch, err := clientset.CoreV1().Pods("").Watch(context.TODO(), metav1.ListOptions{
		FieldSelector: "spec.schedulerName=" + schedulerName,
	})
	if err != nil {
		panic(err.Error())
	}

	fmt.Println("Starting custom scheduler...")

	// 3. 处理事件
	for event := range watch.ResultChan() {
		pod, ok := event.Object.(*v1.Pod)
		if !ok {
			continue
		}

		// 只处理未调度的 Pod
		if pod.Spec.NodeName == "" && pod.Status.Phase == v1.PodPending {
			schedulePod(clientset, pod)
		}
	}
}

func buildConfig() (*rest.Config, error) {
	// 优先使用集群内配置，如果外部运行则使用 kubeconfig
	if config, err := rest.InClusterConfig(); err == nil {
		return config, nil
	}
	return clientcmd.BuildConfigFromFlags("", clientcmd.RecommendedHomeFile)
}

func schedulePod(clientset *kubernetes.Clientset, pod *v1.Pod) {
	// 1. 获取所有节点
	nodes, err := clientset.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
	if err != nil {
		fmt.Printf("Error getting nodes: %v\n", err)
		return
	}

	if len(nodes.Items) == 0 {
		fmt.Println("No nodes available")
		return
	}

	// 2. 随机选择一个节点
	node := nodes.Items[rand.Intn(len(nodes.Items))].Name

	// 3. 绑定 Pod 到节点
	binding := &v1.Binding{
		ObjectMeta: metav1.ObjectMeta{
			Name:      pod.Name,
			Namespace: pod.Namespace,
		},
		Target: v1.ObjectReference{
			APIVersion: "v1",
			Kind:       "Node",
			Name:       node,
		},
	}

	err = clientset.CoreV1().Pods(pod.Namespace).Bind(
		context.TODO(),
		binding,
		metav1.CreateOptions{},
	)

	if err != nil {
		fmt.Printf("Failed to bind pod: %v\n", err)
	} else {
		fmt.Printf("Assigned pod %s/%s to node %s\n", pod.Namespace, pod.Name, node)
	}
}
